# -*- coding: utf-8 -*-

import os
import glob, re, codecs
import logging

from django.conf import settings as config

from basic_info.models import Stock
from companies.models import DailyMoneyFlow, DailyMoneyFlowByIndustry, DailyMoneyFlowByConcept

logger = logging.getLogger(__name__)

class MoneyflowProcessor:
    '''
    计算资金流量模块
    '''
    rt_header_pattern = re.compile(config.MONEYFLOW_DATA_PATTERNS['rt_header'])
    rt_sec_code_pattern = re.compile(config.MONEYFLOW_DATA_PATTERNS['rt_sec_code'])
    rt_sec_name_pattern = re.compile(config.MONEYFLOW_DATA_PATTERNS['rt_sec_name'], re.U)
    rt_date8_pattern = re.compile(config.MONEYFLOW_DATA_PATTERNS['rt_date_8'])
    rt_date4_pattern = re.compile(config.MONEYFLOW_DATA_PATTERNS['rt_date_4'])
    rt_price_pattern = re.compile(config.MONEYFLOW_DATA_PATTERNS['rt_price'])
    rt_reserve_pattern = re.compile(config.MONEYFLOW_DATA_PATTERNS['rt_reserve'])

    def __init__(self):
        self.sec_moneyflow_map = dict()

    def process(self, date_str, root_file_path=None):
        if root_file_path == None:
            root_file_dir = config.MONEYFLOW_SERVICE['daily_file_dir']
        else:
            root_file_dir = root_file_path

        file_name_prefix = config.MONEYFLOW_SERVICE['file_name_prefix']
        date_suffix = date_str.replace('-', '') + "*"
        flow_file_dir = os.path.join(root_file_dir, date_str, file_name_prefix+date_suffix)

        flow_file_list = sorted(glob.glob(flow_file_dir)) # 这里按时间排序非常重要，否则计算的资金流就会错乱

        for flow_file in flow_file_list:
            self.__process_a_file(flow_file)

        item_count = 0
        for sec_code, values in self.sec_moneyflow_map.iteritems():
            date = values[0]
            latest_price = float(values[1])
            total_flow = float(values[2])
            super_flow = float(values[3])
            big_flow = float(values[4])
            non_big_flow = float(values[5])

            try:
                stock = Stock.objects.get(code=sec_code)
            except Stock.DoesNotExist:
                logger.warn("Stock code %s does not exist!", sec_code)
            else:
                try:
                    obj, created = DailyMoneyFlow.objects.get_or_create(date=date_str, 
                                                                        stock=stock,
                                                                        defaults={
                                                                            'closed_price': latest_price,
                                                                            'total_flow': total_flow,
                                                                            'super_flow': super_flow,
                                                                            'big_flow': big_flow,
                                                                            'non_big_flow': non_big_flow,
                                                                        })
                    if not created :
                        obj.closed_price = latest_price
                        obj.total_flow = total_flow
                        obj.super_flow = super_flow
                        obj.big_flow = big_flow
                        obj.non_big_flow = non_big_flow
                        obj.save()

                except Exception as e:
                    # 如果是个别数据解析错误，则忽略该条
                    logger.error(e)

            item_count += 1

        self.__calc_stats_by_concepts(date_str)
        self.__calc_stats_by_industries(date_str)
          
        return item_count
      
    def __process_a_file(self, flow_file_name):
        try:
            with codecs.open(flow_file_name, mode='r', 
                    encoding=config.MONEYFLOW_SERVICE['file_encoding']) as flow_file:
              moneyflow_data = [line.strip() for line in flow_file]
              line_idx = 0
              while line_idx < len(moneyflow_data):
                    line_idx = self.__search_realtime_flow_header(moneyflow_data, line_idx)
        except Exception as e:
            logger.exception('File name: %s', flow_file_name)
            raise e

    def __search_realtime_flow_header(self, money_flow_data, line_start):
        line_idx = line_start
        header = None
        while line_idx < len(money_flow_data):
            header =  self.rt_header_pattern.match(money_flow_data[line_idx])
            if header != None:
                break
            else:
                line_idx += 1

        if header != None:
            # 股票代码
            line_idx += 1
            group_id = int(config.MONEYFLOW_DATA_PATTERN_GROUPS['sec_code_group'])
            sec_code = self.__match_pattern(money_flow_data[line_idx],\
                                 self.rt_sec_code_pattern,\
                                 group_id)
            # 股票名称
            line_idx += 1
            group_id = int(config.MONEYFLOW_DATA_PATTERN_GROUPS['sec_name_group'])
            sec_name = self.__match_pattern(money_flow_data[line_idx],\
                                 self.rt_sec_name_pattern,\
                                 group_id)
            # 交易日期
            line_idx += 1
            group_id = int(config.MONEYFLOW_DATA_PATTERN_GROUPS['date_group'])
            date = self.__match_pattern(money_flow_data[line_idx],\
                                 self.rt_date8_pattern,\
                                 group_id)
            # 流量
            line_idx += 2
            flow_columns = money_flow_data[line_idx].split(config.MONEYFLOW_DATA_DELIMITER)
            # 最新价
            column_id = int(config.MONEYFLOW_DATA_COLUMN_IDS['latest_price_id'])
            group_id = int(config.MONEYFLOW_DATA_PATTERN_GROUPS['price_group'])
            latest_price = self.__match_pattern(flow_columns[column_id],\
                                 self.rt_price_pattern,\
                                 group_id)
            # 总流量
            column_id = int(config.MONEYFLOW_DATA_COLUMN_IDS['total_flow_id'])
            group_id = int(config.MONEYFLOW_DATA_PATTERN_GROUPS['price_group'])
            total_flow = self.__match_pattern(flow_columns[column_id],\
                                 self.rt_price_pattern,\
                                 group_id)
            # 超大流量
            column_id = int(config.MONEYFLOW_DATA_COLUMN_IDS['super_flow_id'])
            group_id = int(config.MONEYFLOW_DATA_PATTERN_GROUPS['price_group'])
            super_flow = self.__match_pattern(flow_columns[column_id],\
                                 self.rt_price_pattern,\
                                 group_id)
            # 大宗流量
            column_id = int(config.MONEYFLOW_DATA_COLUMN_IDS['big_flow_id'])
            group_id = int(config.MONEYFLOW_DATA_PATTERN_GROUPS['price_group'])
            big_flow = self.__match_pattern(flow_columns[column_id],\
                                 self.rt_price_pattern,\
                                 group_id)
            # 非大宗流量
            column_id = int(config.MONEYFLOW_DATA_COLUMN_IDS['non_big_flow_id'])
            group_id = int(config.MONEYFLOW_DATA_PATTERN_GROUPS['price_group'])
            non_big_flow = self.__match_pattern(flow_columns[column_id],\
                                 self.rt_price_pattern,\
                                 group_id)

            self.sec_moneyflow_map[sec_code] = (date, latest_price, total_flow, super_flow, big_flow, non_big_flow)

        return line_idx + 1 # return the next line for following processing

    def __match_pattern(self, line, pattern, group_id):
        try:
            match = pattern.match(line)
            return match.group(group_id)
        except Exception as e:
            logger.exception('Corrupt flow data format at %s.', line)
            raise e

    def __calc_stats_by_industries(self, date_str):
        count = 0
        industry_model = DailyMoneyFlowByIndustry.get_model()
        for industry in industry_model.objects.all():
            sum_total_flow = 0.0
            for stock in industry.stocks.all():
                if self.sec_moneyflow_map.has_key(stock.code):
                    flows = self.sec_moneyflow_map[stock.code]
                    sum_total_flow += float(flows[2]) # total flow
            
            obj, created = DailyMoneyFlowByIndustry.objects.get_or_create(date = date_str,
                                                                          industry=industry,
                                                                          defaults={
                                                                              'total_flow': sum_total_flow,
                                                                          })
            if not created:
                obj.total_flow = sum_total_flow
                obj.save()

            count += 1
        logger.info('%d industries added.', count)

    def __calc_stats_by_concepts(self, date_str):
        count = 0
        concept_model = DailyMoneyFlowByConcept.get_model()
        for concept in concept_model.objects.all():
            sum_total_flow = 0.0
            for stock in concept.stocks.all():
                if self.sec_moneyflow_map.has_key(stock.code):
                    flows = self.sec_moneyflow_map[stock.code]
                    sum_total_flow += float(flows[2]) # total flow
            
            obj, created = DailyMoneyFlowByConcept.objects.get_or_create(date = date_str,
                                                                         concept=concept,
                                                                         defaults={
                                                                             'total_flow': sum_total_flow,
                                                                         })
            if not created:
                obj.total_flow = sum_total_flow
                obj.save()

            count += 1
        logger.info('%d concepts added.', count)

